from PIL import Image
import numpy as np
from utils import *
from dct import *

# 参数
dirpath = 'res/lena.jpg'

# 图像导入
img = np.array(Image.open(dirpath))
print(img.shape)
height, width = img.shape[0:2]
img_num = int(width / 8)

# 图像分块
img_ar = div_pic(img, 8, img_num)
# DCT
dct_img = []
for i in range(img_num):
    dct_img.append(dct_2d(img_ar[i]))

# 创建分布图
img_dis = []
for i in range(8):
    ti = []
    for j in range(8):
        tj = []
        for k in range(len(dct_img)):
            for z in range(len(dct_img[k])):
                tj.append(dct_img[i][j][k][z])
        ti.append(tj)
    img_dis.append(ti)

# 牛顿法求解
for i in range(8):
    for j in range(8):
        if i or j:
            x = np.array(img_dis[i][j])
            abx = np.maximum(x, -x)
            avg_m1 = np.sum(abx) / len(abx)

            x = np.sqrt(x)
            avg_m2 = np.sum(x) / len(abx)

            b = newton_slove(0.54, avg_m1, avg_m2)
            t = np.sum(np.power(abx, b)) / len(abx)
            a = pow(b * t, 1 / b)
        if i == 0 and j == 0:
            print('-', end="\t")
        else:
            print('(' + str(round(a, 2)) + ',' + str(round(b, 2)) + ')', end="\t")
    print('\n')
